import logging
from Crypto.Cipher import AES, PKCS1_v1_5
from Crypto.Random import get_random_bytes
from Crypto.Hash import SHA256, HMAC
from Crypto.Util.Padding import pad
from Crypto.PublicKey import RSA
import base64
import time


class AkamaiCrypto:
    def __init__(self, starttime, platform="android"):
        self.rsakey = None
        self.b = get_random_bytes(16)
        self.c = get_random_bytes(32)
        self.b64aeskey = None
        self.b64hmackey = None
        self.key = None
        self.iv = None
        self.cipher = None
        self.starttime = starttime
        if platform == "android":
            self.platform = "a"
        elif platform == "ios":
            self.platform = "i"

    def import_rsakey(self, key):
        self.key = key
        self.rsakey = RSA.import_key(base64.b64decode(self.key))
        self.cipher = PKCS1_v1_5.new(self.rsakey)
        self.b64aeskey = base64.b64encode(self.cipher.encrypt(self.b))
        logging.debug("AES Key Generated: " + self.b64aeskey.decode("utf-8"))
        self.b64hmackey = base64.b64encode(self.cipher.encrypt(self.c))
        logging.debug("HMAC Key Generated: " + self.b64hmackey.decode("utf-8"))

    def encrypt_message(self, inputstr):
        uptimemillis = time.time_ns() // 1000000 - self.starttime
        self.iv = get_random_bytes(16)
        uptimemillis2 = time.time_ns() // 1000000 - self.starttime
        self.cipher = AES.new(self.b, AES.MODE_CBC, self.iv)
        ct = self.cipher.encrypt(pad(bytes(inputstr, "utf-8"), 16))
        uptimemillis3 = (
            time.time_ns() // 1000000 - self.starttime - uptimemillis2
        ) * 1000
        st = self.iv + ct
        ht = HMAC.new(self.c, st, digestmod=SHA256)
        uptimemillis4 = time.time_ns() // 1000000 - self.starttime
        h = ht.digest()
        uptimemillis5 = (
            time.time_ns() // 1000000 - self.starttime - uptimemillis4
        ) * 1000
        uptimemillis6 = time.time_ns() // 1000000 - self.starttime
        body = base64.b64encode(st + h)
        uptimemillis7 = (
            time.time_ns() // 1000000 - self.starttime - uptimemillis6
        ) * 1000

        outstr = "1,{platform},{b64aeskey},{b64hmackey}${body}${aestime},{hmactime},{b64time}".format(
            platform=self.platform,
            b64aeskey=self.b64aeskey.decode("utf-8"),
            b64hmackey=self.b64hmackey.decode("utf-8"),
            body=body.decode("utf-8"),
            aestime=str(uptimemillis3),
            hmactime=str(uptimemillis5),
            b64time=str(uptimemillis7),
        )
        logging.debug("Encrypted output: " + outstr)
        return outstr
